// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/serializers/presto_serializer_serialization_utils.h>
#include <pollux/vector/base_vector.h>
#include <pollux/vector/constant_vector.h>
#include <pollux/vector/dictionary_vector.h>
#include <pollux/vector/vector_stream.h>
#include <pollux/vector/vector_type_utils.h>

namespace kumo::pollux::serializer::presto::detail {
    template<TypeKind Kind>
    void estimateFlatSerializedSize(
        const BaseVector *vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes) {
        const auto valueSize = vector->type()->cppSizeInBytes();
        if (vector->may_have_nulls()) {
            const auto *rawNulls = vector->raw_nulls();
            for (int32_t i = 0; i < ranges.size(); ++i) {
                const auto end = ranges[i].begin + ranges[i].size;
                const auto numValues = bits::countBits(rawNulls, ranges[i].begin, end);
                // Add the size of the values.
                *(sizes[i]) += numValues * valueSize;
                // Add the size of the null bit mask if there are nulls in the range.
                if (numValues != ranges[i].size) {
                    *(sizes[i]) += bits::nbytes(ranges[i].size);
                }
            }
        } else {
            for (int32_t i = 0; i < ranges.size(); ++i) {
                // Add the size of the values (there's not bit mask since there are no
                // nulls).
                *(sizes[i]) += ranges[i].size * valueSize;
            }
        }
    }

    void estimateFlatSerializedSizeVarcharOrVarbinary(
        const BaseVector *vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes);

    template<>
    inline void estimateFlatSerializedSize<TypeKind::VARCHAR>(
        const BaseVector *vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes) {
        estimateFlatSerializedSizeVarcharOrVarbinary(vector, ranges, sizes);
    }

    template<>
    inline void estimateFlatSerializedSize<TypeKind::VARBINARY>(
        const BaseVector *vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes) {
        estimateFlatSerializedSizeVarcharOrVarbinary(vector, ranges, sizes);
    }

    template<>
    inline void estimateFlatSerializedSize<TypeKind::OPAQUE>(
        const BaseVector *,
        const melon::Range<const IndexRange *> &,
        vector_size_t **) {
        POLLUX_FAIL("Opaque type support is not implemented.");
    }

    void estimateSerializedSizeInt(
        const BaseVector *vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        Scratch &scratch);

    void estimateSerializedSizeInt(
        const BaseVector *vector,
        const melon::Range<const vector_size_t *> &rows,
        vector_size_t **sizes,
        Scratch &scratch);

    void estimateWrapperSerializedSize(
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        const BaseVector *wrapper,
        Scratch &scratch);

    template<TypeKind Kind>
    void estimateDictionarySerializedSize(
        const VectorPtr &vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        Scratch &scratch) {
        POLLUX_CHECK_EQ(vector->encoding(), VectorEncoding::Simple::DICTIONARY);
        using T = typename KindToFlatVector<Kind>::WrapperType;
        auto dictionary_vector = vector->as<DictionaryVector<T> >();

        // We don't currently support serializing DictionaryVectors with nulls, so use
        // the flattened size.
        if (dictionary_vector->nulls()) {
            estimateWrapperSerializedSize(ranges, sizes, vector.get(), scratch);
            return;
        }

        // This will ultimately get passed to simd::transpose, so it needs to be a
        // raw_vector.
        raw_vector<vector_size_t> childIndices;
        std::vector<vector_size_t *> childSizes;
        for (int rangeIndex = 0; rangeIndex < ranges.size(); rangeIndex++) {
            ScratchPtr<vector_size_t, 64> selectedIndicesHolder(scratch);
            auto *mutableSelectedIndices =
                    selectedIndicesHolder.get(dictionary_vector->value_vector()->size());
            auto numUsed = computeSelectedIndices(
                dictionary_vector,
                ranges.subpiece(rangeIndex, 1),
                scratch,
                mutableSelectedIndices);
            for (int i = 0; i < numUsed; i++) {
                childIndices.push_back(mutableSelectedIndices[i]);
                childSizes.push_back(sizes[rangeIndex]);
            }

            // Add the size of the indices.
            *sizes[rangeIndex] += sizeof(int32_t) * ranges[rangeIndex].size;
        }

        // In PrestoBatchVectorSerializer we don't preserve the encodings for the
        // value_vector for a DictionaryVector.
        estimateSerializedSizeInt(
            dictionary_vector->value_vector().get(),
            childIndices,
            childSizes.data(),
            scratch);
    }

    template<TypeKind Kind>
    void estimateConstantSerializedSize(
        const VectorPtr &vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        Scratch &scratch) {
        POLLUX_CHECK(vector->encoding() == VectorEncoding::Simple::CONSTANT);
        using T = typename KindToFlatVector<Kind>::WrapperType;
        auto constant_vector = vector->as<ConstantVector<T> >();
        vector_size_t elementSize = 0;
        if (constant_vector->is_null_at(0)) {
            // There's just a bit mask for the one null.
            elementSize = 1;
        } else if (constant_vector->value_vector()) {
            std::vector<IndexRange> newRanges;
            newRanges.push_back({constant_vector->index(), 1});
            auto *elementSizePtr = &elementSize;
            // In PrestoBatchVectorSerializer we don't preserve the encodings for the
            // value_vector for a ConstantVector.
            estimateSerializedSizeInt(
                constant_vector->value_vector().get(),
                newRanges,
                &elementSizePtr,
                scratch);
        } else if constexpr (std::is_same_v<T, StringView>) {
            elementSize = constant_vector->value_at(0).size();
        } else {
            elementSize = sizeof(T);
        }

        for (int32_t i = 0; i < ranges.size(); ++i) {
            *sizes[i] += elementSize;
        }
    }

    void expandRepeatedRanges(
        const BaseVector *vector,
        const vector_size_t *rawOffsets,
        const vector_size_t *rawSizes,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        std::vector<IndexRange> *childRanges,
        std::vector<vector_size_t *> *childSizes);
} // namespace kumo::pollux::serializer::presto::detail
